/*
 * Copyright (C) 2021 Ignite Realtime Foundation. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.jivesoftware.util;

import org.dom4j.Document;
import org.dom4j.Element;
import org.dom4j.io.SAXReader;
import org.jivesoftware.openfire.JMXManager;
import org.jivesoftware.openfire.mbean.ThreadPoolExecutorDelegate;
import org.jivesoftware.openfire.mbean.ThreadPoolExecutorDelegateMBean;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.xml.sax.SAXException;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.*;
import java.time.Duration;
import java.time.temporal.ChronoUnit;
import java.util.concurrent.*;

/**
 * A utility class that provides a level of abstraction around the creation of a SAXReader instance suitable for XMPP
 * parsing, as well as maintaining a set of readers ready, available for (bulk) processing.
 *
 * @author Guus der Kinderen, guus.der.kinderen@gmail.com
 */
public final class SAXReaderUtil
{
    // Note: this combines a ExecutorService with ThreadLocal usage, which introduces the potential to introduce memory
    //       leaks. The tasks executed by the service make use of a thread-local parser, to avoid the overhead of having
    //       to recreate these. When threads are terminated (after inactivity), this parser cannot be cleaned up (the
    //       API of ThreadPoolExecutor does not provide this functionality). This means that a reference to the value of
    //       the ThreadLocal (the parser) will retain a reference to its ClassLoader, which prevents garbage collection.
    //       As the class of the instance of the parser (SAXReader.class) is loaded by the core classloader in Openfire,
    //       this should only prevent garbage collection of a ClassLoader that should only be garbage collected when
    //       Openfire is shutting down. However, this structure MUST NOT be used using for a value a Class that's loaded
    //       by a plugin. The reference to the PluginClassLoader won't be cleaned, which prevents garbage collection of
    //       the plugin, which prevents the plugin from properly unloading (eg, when it is updated/restarted), which
    //       over time will lead to a wide variety of hard-to-diagnose issues.

    private static final Logger Log = LoggerFactory.getLogger(SAXReaderUtil.class);

    public static final SystemProperty<Integer> PARSER_SERVICE_CORE_POOL_SIZE = SystemProperty.Builder.ofType(Integer.class)
                                                                                    .setKey("xmpp.xmlutil.parser.core-pool-size")
                                                                                    .setMinValue(0)
                                                                                    .setDefaultValue(0)
                                                                                    .setDynamic(false)
                                                                                    .build();

    public static final SystemProperty<Integer> PARSER_SERVICE_MAX_POOL_SIZE = SystemProperty.Builder.ofType(Integer.class)
                                                                                    .setKey("xmpp.xmlutil.parser.maximum-pool-size")
                                                                                    .setMinValue(1)
                                                                                    .setDefaultValue(25)
                                                                                    .setDynamic(false)
                                                                                    .build();

    public static final SystemProperty<Duration> PARSER_SERVICE_KEEP_ALIVE_TIME = SystemProperty.Builder.ofType(Duration.class)
                                                                                    .setKey("xmpp.xmlutil.parser.keep_alive_time")
                                                                                    .setChronoUnit(ChronoUnit.SECONDS)
                                                                                    .setMinValue(Duration.ofMillis(0))
                                                                                    .setDefaultValue(Duration.ofMinutes(1))
                                                                                    .setDynamic(false)
                                                                                    .build();

    private static final ThreadPoolExecutor parserService = new ThreadPoolExecutor(
                                                            PARSER_SERVICE_CORE_POOL_SIZE.getValue(),
                                                            PARSER_SERVICE_MAX_POOL_SIZE.getValue(),
                                                            PARSER_SERVICE_KEEP_ALIVE_TIME.getValue().toMillis(), TimeUnit.MILLISECONDS,
                                                            new SynchronousQueue<>());

    static
    {
        if (JMXManager.isEnabled()) {
            final ThreadPoolExecutorDelegateMBean mBean = new ThreadPoolExecutorDelegate(parserService);
            JMXManager.tryRegister(mBean, ThreadPoolExecutorDelegateMBean.BASE_OBJECT_NAME + "saxReaderUtil");
        }
    }

    /**
     * Schedules parsing of the provided input stream into an XML Document as an asynchronous process. This method
     * returns a {@link Future} object that can be used to obtain the result of the asynchronous process.
     *
     * The XML parsing is delegated to a limited set of parsers, that is fronted by an Executor Service. The amount of
     * parsers that are available is equal to the amount of threads available in the Executor Service. This can be
     * optimized by using the SystemProperties in this class.
     *
     * @param stream The data to be parsed as an XML Document.
     * @return A Future representing pending completion of parsing the provided data as an XML Document.
     */
    public static Future<Document> readDocumentAsync(@Nonnull final InputStream stream) {
        return parserService.submit(new ParserTask(stream));
    }

    /**
     * Parses an XML Document from the content provided by an input stream.
     *
     * Equivalent to calling <tt>readDocumentAsync(input).get()</tt>
     *
     * The XML parsing is delegated to a limited set of parsers, as documented in {@link #readDocumentAsync(InputStream)}.
     * When a parser is not immediately available, this method will block until one becomes available.
     *
     * @param stream The data to be parsed as an XML Document.
     * @return an XML Document.
     * @throws ExecutionException on any error that occurs during parsing (this typically has a {@link org.dom4j.DocumentException} as its cause).
     * @throws InterruptedException when the task that is scheduled to perform the parsing gets interrupted during the execution of the task.
     */
    public static Document readDocument(@Nonnull final InputStream stream) throws ExecutionException, InterruptedException {
        return readDocumentAsync(stream).get();
    }

    /**
     * Parses an XML Document from the content provided by an input stream and returns the root element of the document.
     *
     * Equivalent to calling <tt>readDocument(input).getRootElement()</tt>
     *
     * @param stream The data to be parsed as an XML Document.
     * @return an XML Document root element.
     * @throws ExecutionException on any error that occurs during parsing (this typically has a {@link org.dom4j.DocumentException} as its cause).
     * @throws InterruptedException when the task that is scheduled to perform the parsing gets interrupted during the execution of the task.
     */
    public static Element readRootElement(@Nonnull final InputStream stream) throws ExecutionException, InterruptedException {
        return readDocument(stream).getRootElement();
    }

    /**
     * Schedules parsing of the provided text into an XML Document as an asynchronous process. This method
     * returns a {@link Future} object that can be used to obtain the result of the asynchronous process.
     *
     * The XML parsing is delegated to a limited set of parsers, that is fronted by an Executor Service. The amount of
     * parsers that are available is equal to the amount of threads available in the Executor Service. This can be
     * optimized by using the SystemProperties in this class.
     *
     * @param text The data to be parsed as an XML Document.
     * @return A Future representing pending completion of parsing the provided data as an XML Document.
     */
    public static Future<Document> readDocumentAsync(@Nonnull final String text) {
        return parserService.submit(new ParserTask(new StringReader(text)));
    }

    /**
     * Parses an XML Document from the provided text.
     *
     * Equivalent to calling <tt>readDocumentAsync(input).get()</tt>
     *
     * The XML parsing is delegated to a limited set of parsers, as documented in {@link #readDocumentAsync(String)}.
     * When a parser is not immediately available, this method will block until one becomes available.
     *
     * @param text The data to be parsed as an XML Document.
     * @return an XML Document.
     * @throws ExecutionException on any error that occurs during parsing (this typically has a {@link org.dom4j.DocumentException} as its cause).
     * @throws InterruptedException when the task that is scheduled to perform the parsing gets interrupted during the execution of the task.
     */
    public static Document readDocument(@Nonnull final String text) throws ExecutionException, InterruptedException {
        return readDocumentAsync(text).get();
    }

    /**
     * Parses an XML Document from the provided text and returns the root element of the document.
     *
     * Equivalent to calling <tt>readDocument(input).getRootElement()</tt>
     *
     * @param text The data to be parsed as an XML Document.
     * @return an XML Document root element.
     * @throws ExecutionException on any error that occurs during parsing (this typically has a {@link org.dom4j.DocumentException} as its cause).
     * @throws InterruptedException when the task that is scheduled to perform the parsing gets interrupted during the execution of the task.
     */
    public static Element readRootElement(@Nonnull final String text) throws ExecutionException, InterruptedException {
        return readDocument(text).getRootElement();
    }

    /**
     * Schedules parsing of the provided reader into an XML Document as an asynchronous process. This method
     * returns a {@link Future} object that can be used to obtain the result of the asynchronous process.
     *
     * The XML parsing is delegated to a limited set of parsers, that is fronted by an Executor Service. The amount of
     * parsers that are available is equal to the amount of threads available in the Executor Service. This can be
     * optimized by using the SystemProperties in this class.
     *
     * @param reader The data to be parsed as an XML Document.
     * @return A Future representing pending completion of parsing the provided data as an XML Document.
     */
    public static Future<Document> readDocumentAsync(@Nonnull final Reader reader) {
        return parserService.submit(new ParserTask(reader));
    }

    /**
     * Parses an XML Document from the content provided by a reader.
     *
     * Equivalent to calling <tt>readDocumentAsync(input).get()</tt>
     *
     * The XML parsing is delegated to a limited set of parsers, as documented in {@link #readDocumentAsync(Reader)}.
     * When a parser is not immediately available, this method will block until one becomes available.
     *
     * @param reader The data to be parsed as an XML Document.
     * @return an XML Document.
     * @throws ExecutionException on any error that occurs during parsing (this typically has a {@link org.dom4j.DocumentException} as its cause).
     * @throws InterruptedException when the task that is scheduled to perform the parsing gets interrupted during the execution of the task.
     */
    public static Document readDocument(@Nonnull final Reader reader) throws ExecutionException, InterruptedException {
        return readDocumentAsync(reader).get();
    }

    /**
     * Parses an XML Document from the content provided by a reader and returns the root element of the document.
     *
     * Equivalent to calling <tt>readDocument(input).getRootElement()</tt>
     *
     * @param reader The data to be parsed as an XML Document.
     * @return an XML Document root element.
     * @throws ExecutionException on any error that occurs during parsing (this typically has a {@link org.dom4j.DocumentException} as its cause).
     * @throws InterruptedException when the task that is scheduled to perform the parsing gets interrupted during the execution of the task.
     */
    public static Element readRootElement(@Nonnull final Reader reader) throws ExecutionException, InterruptedException {
        return readDocument(reader).getRootElement();
    }

    /**
     * Schedules parsing of content of afile into an XML Document as an asynchronous process. This method
     * returns a {@link Future} object that can be used to obtain the result of the asynchronous process.
     *
     * The XML parsing is delegated to a limited set of parsers, that is fronted by an Executor Service. The amount of
     * parsers that are available is equal to the amount of threads available in the Executor Service. This can be
     * optimized by using the SystemProperties in this class.
     *
     * @param file The data to be parsed as an XML Document.
     * @return A Future representing pending completion of parsing the provided data as an XML Document.
     */
    public static Future<Document> readDocumentAsync(@Nonnull final File file) {
        return parserService.submit(new ParserTask(file));
    }

    /**
     * Parses an XML Document from the content provided by a file.
     *
     * Equivalent to calling <tt>readDocumentAsync(input).get()</tt>
     *
     * The XML parsing is delegated to a limited set of parsers, as documented in {@link #readDocumentAsync(File)}.
     * When a parser is not immediately available, this method will block until one becomes available.
     *
     * @param file The data to be parsed as an XML Document.
     * @return an XML Document.
     * @throws ExecutionException on any error that occurs during parsing (this typically has a {@link org.dom4j.DocumentException} as its cause).
     * @throws InterruptedException when the task that is scheduled to perform the parsing gets interrupted during the execution of the task.
     */
    public static Document readDocument(@Nonnull final File file) throws ExecutionException, InterruptedException {
        return readDocumentAsync(file).get();
    }

    /**
     * Parses an XML Document from the content provided by a file and returns the root element of the document.
     *
     * Equivalent to calling <tt>readDocument(input).getRootElement()</tt>
     *
     * @param file The data to be parsed as an XML Document.
     * @return an XML Document root element.
     * @throws ExecutionException on any error that occurs during parsing (this typically has a {@link org.dom4j.DocumentException} as its cause).
     * @throws InterruptedException when the task that is scheduled to perform the parsing gets interrupted during the execution of the task.
     */
    public static Element readRootElement(@Nonnull final File file) throws ExecutionException, InterruptedException {
        return readDocument(file).getRootElement();
    }

    /**
     * Constructs a new SAXReader instance, configured for XMPP parsing.
     *
     * @return A new SAXReader instance.
     */
    private static SAXReader constructNewReader() throws SAXException
    {
        final SAXReader saxReader = new SAXReader();
        saxReader.setEncoding( "UTF-8" );
        saxReader.setIgnoreComments(true);

        // Guard against various known vulnerabilities (see OF-2094).
        saxReader.setEntityResolver((publicId, systemId) -> {
            throw new IOException("External entity denied: " + publicId + " // " + systemId);
        });
        saxReader.setFeature("http://apache.org/xml/features/nonvalidating/load-external-dtd", false);
        saxReader.setFeature("http://apache.org/xml/features/disallow-doctype-decl", true);
        saxReader.setFeature("http://xml.org/sax/features/external-general-entities", false);
        saxReader.setFeature("http://xml.org/sax/features/external-parameter-entities", false);

        return saxReader;
    }

    private static class ParserTask implements Callable<Document>
    {
        private static final ThreadLocal<SAXReader> localSAXReader = ThreadLocal.withInitial(()-> {
            try {
                return constructNewReader();
            } catch (SAXException e) {
                Log.error("Unable to construct a new XML parser.", e);
                return null;
            }
        });

        @Nullable
        private final InputStream stream;

        @Nullable
        private final Reader reader;

        @Nullable
        private final File file;

        public ParserTask(@Nonnull final InputStream stream) {
            this.stream = stream;
            this.reader = null;
            this.file = null;
        }

        public ParserTask(@Nonnull final Reader reader) {
            this.stream = null;
            this.reader = reader;
            this.file = null;
        }

        public ParserTask(@Nonnull final File file) {
            this.stream = null;
            this.reader = null;
            this.file = file;
        }

        @Override
        public Document call() throws Exception {
            if (stream != null) {
                return localSAXReader.get().read(stream);
            } else if (reader != null) {
                return localSAXReader.get().read(reader);
            } else if (file != null) {
                return localSAXReader.get().read(file);
            } else {
                // Did you forget to add an 'else' block for the new type that you added?
                throw new IllegalStateException("Unable to parse data. Data is either null, or of an unrecognized type.");
            }
        }
    }

//    // This allows for a quick and dirty verification that the SAXReader instances that are values of the ThreadLocal
//    // are indeed garbage collected after the threads expire. To verify, use jmap to dump a histogram that shows the
//    // number of instances directly after starting the application, and after the first 'sleep' has expired (and all
//    // threads should have been terminated by the Service).
//    public static void main(String[] args) throws ExecutionException, InterruptedException {
//        for (int i = 0; i<300; i++) {
//            XmlUtils.readDocument(new java.io.StringReader("<foo></foo>"));
//        }
//        System.out.println("Histo should now have up to 25 SAXReader instances.");
//        Thread.sleep(65_000); // wait for threads to time out
//        System.gc();
//        System.out.println("Histo should now have 0 SAXReader instances.");
//        Thread.sleep(65_000);
//    }
}
